import cv2
import easyocr
import numpy as np
import pyautogui
from PIL import Image
from pathlib import Path

root_dir = Path(__file__).resolve().parent.parent.parent

# 初始化识别器（中英文）
custom_model_dir = root_dir / "my_ocr_models"


def img_ocr(region=None):
    # 创建目录（如果不存在）
    custom_model_dir.mkdir(parents=True, exist_ok=True)
    reader = easyocr.Reader(
        lang_list=['ch_sim', 'en'],
        download_enabled=True,
        model_storage_directory=str(custom_model_dir),
        user_network_directory=str(custom_model_dir / 'user_networks'),
    )
    screenshot = pyautogui.screenshot(region=(
        region[0], 
        region[1], 
        region[2], 
        region[3]
    ))
            # 将PIL图像转为OpenCV格式
    img = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)

    red_channel = img[:, :, 2].astype(np.float32)
    red_enhanced = cv2.normalize(red_channel*2.0, None, 0, 255, cv2.NORM_MINMAX) 
    red_gray = cv2.normalize(red_enhanced, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
    
    cv2.imwrite("test/red_gray.png", red_gray)

    results = reader.readtext(red_gray, detail=1)
    numbers = []
    for i, (bbox, text, confidence) in enumerate(results):
        # 新增数字提取逻辑
        try:
            # 去除中文和空格后转换数字
            num_str = ''.join(filter(str.isdigit, text))
            number = int(num_str)
            if confidence >0.2:
                numbers.append(number)
            print(f"{i+1}. 内容：'{text}' | 提取数字：{number} | 置信度：{confidence:.2%}")
        except (ValueError, IndexError) as e:
            print(f"{i+1}. 内容：'{text}' | 数字提取失败: {str(e)}")

    return numbers  


if __name__ == "__main__":
    import time
    time.sleep(3)
    num=img_ocr(region=(1095,1418,1259-1095,1470-1418,)) 
    print(num)
